'''
本文件下包含所有计算特征的函数,所有函数统一接口，一个是对应股票的每日收盘价
一个是需要计算的参数
'''
import numpy as np
import pandas as pd

#计算最近站上MA均线的天数,参数为prices，和ma天数的列表
from Utils.tools import findBottoms,findPeaks

def calLargerThanMas(prices,code,madays,prefix):
    theprice=prices['close']
    tmpres={}
    for day in madays:
        ma = pd.DataFrame(theprice.rolling(window=day).mean())
        ma.columns = ['close_ma']
        # print(ma)
        merge = pd.concat([ma, theprice], axis=1)
        merge = merge.dropna()
        merge['diff'] = merge['close'] - merge['close_ma']
        count=0
        data=list(merge['diff'])
        data.reverse()
        for i in range(len(data)-1):
            if data[i]>=0:
                count=count+1
            else:
                break
        tmpres[prefix + 'ma' + str(day)] = [count]
    resdf=pd.DataFrame(tmpres,index=[code])
    return resdf

#计算收盘价是否升穿均线
def closeUpPenetrateMas(prices,code,madays,prefix):
    theprice = prices['close']
    tmpres = {}
    for maday in madays:
        ma=theprice.rolling(window=maday).mean()
        last_tradeday_ma = list(ma.tail(1))[0]
        last_2tradeday_ma = list(ma.tail(2).head(1))[0]
        last_trade_day = list(prices['close'].tail(1))[0]
        last_2trade_day = list(prices['close'].tail(2).head(1))[0]
        if last_2trade_day<last_2tradeday_ma and last_trade_day>=last_tradeday_ma:
            tmpres[prefix+str(maday)]=[True]
        else:
            tmpres[prefix + str(maday)] = [False]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算标准差
def calSTDS(prices,code,stddays,column,prefix):
    theprice = prices[column]
    tmpres = {}
    for day in stddays:
        tmpres[prefix+str(day)]=round(np.std(list(theprice.tail(day))),4)
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf


#计算收盘价低于MA均线的天数
def calLessThanDays(prices,code,madays,prefix):
    tmpres = {}
    prices=prices.head(prices.shape[0]-1)
    price2=prices['close'].dropna()
    shape = prices.shape[0]
    price2=list(price2.head(shape-1))
    price2.reverse()
    for maday in madays:
        count = 0
        ma2 = list(prices['close'].rolling(window=maday).mean().head(shape - 1))
        ma2.reverse()
        for i in range(len(ma2)):
            if ma2[i]>=price2[i]:
                count=count+1
            else:
                break
        tmpres[prefix + str(maday)] = [count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算股价分位数
def calQuantile(prices,code,quantileDays,prefix):
    tmpres = {}
    # current_price = list(prices['close'].tail(1))[0]
    for quantileDay in quantileDays:
        previous_prices=prices.tail(quantileDay)
        # previous_prices.sort()
        # count=0
        # for i in range(len(previous_prices)):
        #     if current_price>=previous_prices[i]:
        #         count=count+1
        #     else:
        #         break
        quant=len(previous_prices[previous_prices['close']<previous_prices.iloc[-1]['close']])/len(previous_prices)
        tmpres[prefix + str(quantileDay)] = [round(quant, 2)]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算涨跌幅
def calIncrease(prices,code,increaseDays,prefix):
    tmpres = {}
    current_price = list(prices['close'].tail(1))[0]
    for increaseDay in increaseDays:
        previous_price=list(prices['close'].tail(increaseDay).head(1))[0]
        increase=0
        if previous_price is np.nan or previous_price == 0:
            increase=np.nan
        else:
            increase = (current_price - previous_price) / previous_price * 100
        tmpres[prefix+str(increaseDay)]=[round(increase,2)]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#MA多头排列
def getMaOrder(prices,code,maDays,prefix):
    tmpres = {}
    mas=[]
    for maday in maDays:
        thema=prices["close"].rolling(window=maday).mean()
        mas.append(thema.tail(1)[0])
    maflag=True
    for i in range(len(maDays)-1):
            if mas[i]<=mas[i+1]:
                maflag=False
                break
    tmpres[prefix+'order']=[maflag]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算均线值
def calMas(prices,code,maDays,prefix):
    tmpres = {}
    for maday in maDays:
        thema = prices["close"].rolling(window=maday).mean()
        lastday_ma=list(thema.tail(1))[0]
        tmpres[prefix+str(maday)]=[lastday_ma]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#MA连续增长的天数
def maIncreaseDays(prices,code,maDays,prefix):
    theprice=prices['close']
    tmpres = {}
    for maday in maDays:
        ma = list(theprice.rolling(window=maday).mean())
        ma.reverse()
        count=0
        for i in range(len(ma)-1):
            if ma[i]>=ma[i+1]:
                count=count+1
            else:
                break
        tmpres[prefix + str(maday)] = [count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#股价连续增长的天数
def closeIncreaseDays(prices,code,prefix):
    theprice = list(prices['close'].dropna())
    theprice.reverse()
    tmpres = {}
    count = 0
    for i in range(len(theprice) - 1):
        if theprice[i] >= theprice[i + 1]:
            count = count + 1
        else:
            break
    tmpres[prefix + 'cnt'] = [count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#股价连续跌的天数
def closeDecreaseDays(prices,code,prefix):
    theprice = list(prices['close'].dropna())
    theprice.reverse()
    tmpres = {}
    count = 0
    for i in range(len(theprice) - 1):
        if theprice[i] < theprice[i + 1]:
            count = count + 1
        else:
            break
    tmpres[prefix + 'cnt'] = [count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算均线之间互相升穿
def masUpPenetrate(prices,code,madaylow,madayhigh,prefix):
    theprice = prices['close']
    tmpres = {}
    malow = theprice.rolling(window=madaylow).mean()
    last_tradeday_malow = list(malow.tail(1))[0]
    last_2tradeday_malow = list(malow.tail(2).head(1))[0]
    mahigh= theprice.rolling(window=madayhigh).mean()
    last_tradeday_mahigh = list(mahigh.tail(1))[0]
    last_2tradeday_mahigh = list(mahigh.tail(2).head(1))[0]
    if last_tradeday_mahigh<=last_tradeday_malow and last_2tradeday_malow<=last_2tradeday_mahigh:
        tmpres[prefix+str(madaylow)+'up'+str(madayhigh)]=[True]
    else:
        tmpres[prefix + str(madaylow) + 'up' + str(madayhigh)] = [False]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算量比
def volRatio(prices,code,prefix):
    thevolume = prices['volume']
    thevolume=thevolume.dropna()
    tmpres = {}
    current_vol = list(thevolume.tail(1))[0]
    yesterday_vol=list(thevolume.tail(2).head(1))[0]
    if yesterday_vol is not np.nan and yesterday_vol is not None and yesterday_vol!=0:
        tmpres[prefix]=[current_vol/yesterday_vol]
    else:
        tmpres[prefix] = [np.nan]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

def volSumRatio(prices,code,prefix):
    thevolume = prices['volume']
    thevolume=thevolume.dropna()
    tmpres = {}
    current_vol = thevolume.tail(5).sum()
    yesterday_vol=thevolume.tail(10).head(5).sum()
    if yesterday_vol is not np.nan and yesterday_vol is not None and yesterday_vol!=0:
        tmpres[prefix]=[current_vol/yesterday_vol]
    else:
        tmpres[prefix] = [np.nan]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算成交量增长的天数
def volIncreaseDays(prices,code,prefix):
    thevolume = prices['volume']
    thevolume=thevolume.dropna()
    tmpres = {}
    volumes=list(thevolume)
    volumes.reverse()
    count=0
    for i in range(len(volumes)-1):
        if volumes[i]>=volumes[i+1]:
            count=count+1
        else:
            break
    tmpres[prefix] = [count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算最近一段时间内相对最高点的跌幅
def decreaseFromMaxRecent(prices,code,decreaseDays,prefix):
    last_trade_day = list(prices['close'].tail(1))[0]
    # theprice = prices['high'].dropna()
    tmpres = {}
    for day in decreaseDays:
        peaks=findPeaks(prices.tail(day)['close'].dropna())
        if len(peaks)>0:
            maxprice=max(peaks)
            if maxprice is not np.nan and maxprice is not None and maxprice!=0:
                tmpres[prefix+str(day)]=[(last_trade_day-maxprice)/maxprice*100]
            else:
                tmpres[prefix + str(day)] = [np.nan]
        else:
            tmpres[prefix + str(day)] = [np.nan]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#从最近底部起来的涨幅
def increaseFromMinRecent(prices,code,increaseDays,prefix):
    last_trade_day = list(prices['close'].tail(1))[0]
    tmpres = {}
    for day in increaseDays:
        bottoms=findBottoms(prices.tail(day)['close'].dropna())
        if len(bottoms)>0:
            minprice=min(bottoms)
            if minprice is not np.nan and minprice is not None and minprice!=0:
                tmpres[prefix+str(day)]=[(last_trade_day-minprice)/minprice*100]
            else:
                tmpres[prefix + str(day)] = np.nan
        else:
            tmpres[prefix + str(day)] = np.nan
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算振幅
def calAmplitude(prices,code,ampDays,prefix):
    tmpres = {}
    high=prices['high'].dropna()
    low=prices['low'].dropna()
    last_trade_day = list(prices['close'].dropna().tail(1))[0]
    for ampday in ampDays:
        highmax=high.tail(ampday).max()
        maxid=high.tail(ampday).idxmax()
        lowmin=low.tail(ampday).min()
        minid=low.tail(ampday).idxmin()
        if last_trade_day is not np.nan and last_trade_day!=0 and last_trade_day is not None:
            if maxid<minid:
                tmpres[prefix+str(ampday)]=[(highmax-lowmin)/highmax*100]
            else:
                tmpres[prefix + str(ampday)] = [(lowmin-highmax) / lowmin * 100]
        else:
            tmpres[prefix + str(ampday)]=np.nan
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算最近一段时间的minmax的比例
def calMinMaxRatio(prices,code,minmaxDays,prefix):
    tmpres = {}
    high=prices['high'].dropna()
    low=prices['low'].dropna()
    last_trade_day = list(prices['close'].dropna().tail(1))[0]
    for minmaxDay in minmaxDays:
        highmax=high.tail(minmaxDay).max()
        lowmin=low.tail(minmaxDay).min()
        if last_trade_day is not np.nan and last_trade_day!=0 and last_trade_day is not None:
            tmpres[prefix+str(minmaxDay)]=[highmax/lowmin]
        else:
            tmpres[prefix + str(minmaxDay)]=np.nan
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#是否最近的新低
def isRecentDaysMin(prices,code,checkDays,prefix):
    tmpres = {}
    last_trade_day = list(prices['close'].dropna().tail(1))[0]
    for checkDay in checkDays:
        theprice = prices['close'].dropna().tail(checkDay)
        minprice = theprice.min()
        if minprice>=last_trade_day:
            tmpres[prefix+str(checkDay)]=[True]
        else:
            tmpres[prefix + str(checkDay)] = [False]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#是否最近的新高
def isRecentDaysMax(prices,code,checkDays,prefix):
    tmpres = {}
    last_trade_day = list(prices['close'].dropna().tail(1))[0]
    for checkDay in checkDays:
        theprice=prices['close'].dropna().tail(checkDay)
        maxprice=theprice.max()
        if maxprice==last_trade_day:
            tmpres[prefix+str(checkDay)]=[True]
        else:
            tmpres[prefix + str(checkDay)] = [False]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#最近一段时间的股价最低
def recentDaysMin(prices,code,checkDays,prefix):
    tmpres = {}
    last_trade_day = list(prices['close'].dropna().tail(1))[0]
    for checkDay in checkDays:
        theprice=prices['close'].dropna().tail(checkDay)
        minprice=theprice.min()
        tmpres[prefix+str(checkDay)]=[minprice]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#最近一段时间的股价最高
def recentDaysMax(prices,code,checkDays,prefix):
    tmpres = {}
    last_trade_day = list(prices['close'].dropna().tail(1))[0]
    for checkDay in checkDays:
        theprice=prices['close'].dropna().tail(checkDay)
        maxprice=theprice.max()
        tmpres[prefix+str(checkDay)]=[maxprice]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf


#获取收盘价df
def getCloseDF(prices,code):
    tmpres = {}
    last_trade_day = list(prices['close'].dropna().tail(1))[0]
    tmpres['close'] = [last_trade_day]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算最近在250日线附近的穿插次数
def calAroundMA(prices,code,recentDays,prefix):
    theprice=prices['close']
    tmpres={}
    ma = pd.DataFrame(theprice.rolling(window=250).mean())
    ma.columns = ['close_ma']
    # print(ma)
    merge = pd.concat([ma, theprice], axis=1)
    merge = merge.dropna()
    merge['diff'] = merge['close'] - merge['close_ma']
    merge=merge.tail(recentDays)
    count=0
    data=list(merge['diff'])
    data.reverse()
    for i in range(len(data)-1):
        if (data[i]>=0 and data[i+1]<=0) or (data[i]<=0 and data[i+1]>=0):
            count=count+1
    tmpres[prefix] = [count]
    resdf=pd.DataFrame(tmpres,index=[code])
    return resdf

#计算某段时间阳线的天数
def calSunDays(prices,code,recentDays,prefix):
    import warnings
    warnings.filterwarnings("ignore")
    tmpres = {}
    for day in recentDays:
        theprice=prices.tail(day)
        theprice['issun']=theprice.apply(lambda x:1 if x['close']>x['open'] else 0,axis=1)
        sundays=theprice[theprice['issun']==1].shape[0]
        tmpres[prefix+str(day)]=[sundays]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算最近一段时间阳线的比例
def calSunDaysRatio(prices,code,recentDays,prefix):
    import warnings
    warnings.filterwarnings("ignore")
    tmpres = {}
    for day in recentDays:
        theprice = prices.tail(day)
        theprice['issun'] = theprice.apply(lambda x: 1 if x['close'] > x['open'] else 0, axis=1)
        sundays = theprice[theprice['issun'] == 1].shape[0]
        tmpres[prefix + str(day)] = [sundays/theprice.shape[0]]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算最近连续阳线的天数
def calRecentSunDays(prices,code,prefix):
    import warnings
    warnings.filterwarnings("ignore")
    tmpres = {}
    prices['issun'] = prices.apply(lambda x: 1 if x['close'] >= x['open'] else 0, axis=1)
    issun=list(prices['issun'])
    issun.reverse()
    count=0
    for i in range(len(issun)):
        if issun[i]==1:
            count=count+1
        else:
            break
    tmpres[prefix]=[count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算最近连续阳线的涨幅
def calRecentSunDaysIncrease(prices,code,prefix):
    import warnings
    warnings.filterwarnings("ignore")
    tmpres = {}
    prices['issun'] = prices.apply(lambda x: 1 if x['close'] >= x['open'] else 0, axis=1)
    issun=list(prices['issun'])
    issun.reverse()
    count=0
    for i in range(len(issun)):
        if issun[i]==1:
            count=count+1
        else:
            break
    inc = 0
    if count > 0:
        current_price = list(prices['close'].tail(1))[0]
        previous_price = list(prices['close'].tail(count).head(1))[0]
        inc = round((current_price-previous_price) / previous_price * 100.0,3)
    tmpres[prefix + 'increase'] = [inc]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算最近阴线的天数
def calRecentNegativeDays(prices,code,prefix):
    import warnings
    warnings.filterwarnings("ignore")
    tmpres = {}
    prices['issun'] = prices.apply(lambda x: 1 if x['close'] >= x['open'] else 0, axis=1)
    issun=list(prices['issun'])
    issun.reverse()
    count=0
    for i in range(len(issun)):
        if issun[i]==0:
            count=count+1
        else:
            break
    tmpres[prefix]=[count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算连续阴线的跌幅
def calRecentNegativeDaysDecrease(prices,code,prefix):
    import warnings
    warnings.filterwarnings("ignore")
    tmpres = {}
    prices['issun'] = prices.apply(lambda x: 1 if x['close'] >= x['open'] else 0, axis=1)
    issun=list(prices['issun'])
    issun.reverse()
    count=0
    for i in range(len(issun)):
        if issun[i]==0:
            count=count+1
        else:
            break
    dec=0
    if count>0:
        current_price=list(prices['close'].tail(1))[0]
        previous_price = list(prices['close'].tail(count).head(1))[0]
        dec=round((previous_price-current_price)/previous_price*100.0,3)
    tmpres[prefix+'decrease']=[dec]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#升穿250均线之后的涨幅
def closeUpPenetrateMa250(prices,code,prefix):
    theprice = prices['close']
    tmpres = {}
    day=250
    ma = pd.DataFrame(theprice.rolling(window=day).mean())
    ma.columns = ['close_ma']
    # print(ma)
    merge = pd.concat([ma, theprice], axis=1)
    merge = merge.dropna()
    merge['diff'] = merge['close'] - merge['close_ma']
    count = 0
    data = list(merge['diff'])
    data.reverse()
    #计算升穿的时间
    for i in range(len(data) - 1):
        if data[i] >= 0:
            count = count + 1
        else:
            break
    inc=0
    if count>0:
        current_price = list(prices['close'].tail(1))[0]
        previous_price = list(prices['close'].tail(count).head(1))[0]
        inc = round((current_price-previous_price) / previous_price * 100.0, 3)
    tmpres[prefix] = [inc]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算HeiKin Ashi曲线
def calHA(prices):
    import warnings
    warnings.filterwarnings("ignore")
    df=prices
    df['ha_close'] = (df.close + df.open + df.high + df.low) / 4.0
    ha_open = np.zeros(df.shape[0])
    ha_open[0] = (df.open[0]+df.close[0])/2.0
    for i in range(1, df.shape[0]):
        ha_open[i] = (ha_open[i - 1] + df['ha_close'][i - 1]) / 2.0
    # df.insert(1, 'ha_open', ha_open)
    df['ha_open']=ha_open
    df['ha_high'] = df[['high', 'ha_open', 'ha_close']].max(axis=1)
    df['ha_low'] = df[['low', 'ha_open', 'ha_close']].min(axis=1)
    df = df.iloc[1:]
    return df

#计算HA连续阴线的天数
def calHANegDays(prices,code,prefix):
    import warnings
    warnings.filterwarnings("ignore")
    tmpres = {}
    prices=calHA(prices)
    prices['issun'] = prices.apply(lambda x: 1 if x['ha_close'] >= x['ha_open'] else 0, axis=1)
    issun = list(prices['issun'])
    issun.reverse()
    count = 0
    for i in range(len(issun)):
        if issun[i] == 0:
            count = count + 1
        else:
            break
    tmpres[prefix] = [count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算HA连续阳线的天数
def calHASunDays(prices,code,prefix):
    import warnings
    warnings.filterwarnings("ignore")
    tmpres = {}
    prices=calHA(prices)
    prices['issun'] = prices.apply(lambda x: 1 if x['ha_close'] >= x['ha_open'] else 0, axis=1)
    issun = list(prices['issun'])
    issun.reverse()
    count = 0
    for i in range(len(issun)):
        if issun[i] == 1:
            count = count + 1
        else:
            break
    tmpres[prefix] = [count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#HA反转股
def HAInvert(prices,code,prefix):
    import warnings
    warnings.filterwarnings("ignore")
    tmpres = {}
    prices = calHA(prices)
    prices['issun'] = prices.apply(lambda x: 1 if x['ha_close'] >= x['ha_open'] else 0, axis=1)
    issun = list(prices['issun'])
    issun.reverse()
    count = 0
    if issun[0]==1:
        #计算趋势反转
        for i in range(1,len(issun)):
            if issun[i] == 0:
                count = count + 1
            else:
                break
    tmpres[prefix] = [count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf

#计算连续跌的拐点
def closeInvert(prices,code,prefix):
    import warnings
    warnings.filterwarnings("ignore")
    tmpres = {}
    prices['issun'] = prices.apply(lambda x: 1 if x['close'] >= x['open'] else 0, axis=1)
    issun = list(prices['issun'])
    issun.reverse()
    count = 0
    if issun[0] == 1:
        # 计算趋势反转
        for i in range(1, len(issun)):
            if issun[i] == 0:
                count = count + 1
            else:
                break
    tmpres[prefix] = [count]
    resdf = pd.DataFrame(tmpres, index=[code])
    return resdf